import os.path

import torch
from transformers import BertTokenizerFast
from transformers import GPT2LMHeadModel
from parameter_config import ParameterConfig
import torch.nn.functional as F
params = ParameterConfig()

PAD = '[PAD]'

pad_id = 0
device = params.device
print(f'当前使用的设备：{device}')

tokenizer = BertTokenizerFast(vocab_file=params.vocab_path,
                                  sep_token='[SEP]',
                                  cls_token='[CLS]',
                                  pad_token='[PAD]')

model = GPT2LMHeadModel.from_pretrained("D:\\Project\\GPT2\\save_model\\binzi_model_epoch1")

model = model.to(device)
model.eval()

def top_k_top_p_filtering(logits, top_k, filter_value=-float('Inf')):
    """
    使用top-k/或者top-p筛选过滤logits的分布
        top-k > 0 : 保留概率最高的top k个标记
    """
    assert logits.dim() == 1
    top_k = min(top_k, logits.size(-1))  # 确保top_k不超过logits的最后一个维度大小
    # 判断topk
    if top_k > 0:
        # indices_to_remove获取的是一个布尔张量列表,用于获取logits中概率值小于top_k的标记
        # 小于topk值默认设置为False,大于是True 示例: [ True, False,  True, False]
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        # 对于topk之外的将logits值设置为负无穷   示例: [  -inf, 3.6000,   -inf, 4.8000]
        logits[indices_to_remove] = filter_value
    return logits

def model_predict(input_text):
    """

    :param input_text: 问题
    :return: 答案
    """
    sample_file = ''
    if params.save_samples_path:
        if not os.path.exists(params.save_samples_path):
            os.makedirs(params.save_samples_path)

        sample_file = open(params.save_samples_path + '/my_model_chat.txt', 'a', encoding='utf-8')

    sample_file.write(f"用户的问题：{input_text}\n")

    history = []
    text_ids = tokenizer.encode(input_text, add_special_tokens=False)
    history.append(text_ids)
    input_ids = [tokenizer.cls_token_id]
    for history_id, history_utr in enumerate(history[-params.max_history_len:]):
        input_ids.extend(history_utr)
        input_ids.append(tokenizer.sep_token_id)

    input_ids = torch.tensor(input_ids, dtype=torch.long, device=device)
    input_ids = input_ids.unsqueeze(dim=0)

    response = []
    for _ in range(params.max_len):
        # TODO 调用模型  也可以用model(input_ids=input_ids)
        outputs = model.forward(input_ids=input_ids)
        # TODO 将模型输出的logits值赋值给变量logits
        logits = outputs.logits
        # TODO 获取下一个单词的概率值
        next_token_logits = logits[0, -1, :]
        # TODO 对reponse已经生成的结果中每个token添加一个重复惩罚,防止模型重复输出已经生成的结果
        for id in set(response):
            next_token_logits[id] /= params.repetition_penalty
        # 对于[UNK]的概率设置为无穷小,也就是说模型的输出中不能出现[UNK]
        next_token_logits[tokenizer.convert_tokens_to_ids('[UNK]')] = -float('Inf')
        # TODO 自定义一个函数,使用top-k和top-p过滤策略,对logits进行过滤,提升生成质量
        # 示例结果: [  -inf, 3.6000,   -inf, 4.8000]
        filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=params.topk)
        # 先用torch中multinomial表示从候选的集合中无放回的抽取num_sample个元素
        # 权重越高,抽到的几率越高,返回元素的下标
        next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1).to(device)
        # 如果下一个标记遇到[SEP]则表名response生成结束,停止循环
        if next_token.item() == tokenizer.sep_token_id:
            break
        # 将下一个标记的id添加到响应列表中
        response.append(next_token.item())
        # TODO 沿着列维度将下一个标记与输入id拼接,更新输入序列
        input_ids = torch.cat((input_ids, next_token.unsqueeze(0)), dim=1)

    # 将机器人的响应添加到对话历史中
    history.append(response)
    # 将响应从id序列转换文本
    res = tokenizer.convert_ids_to_tokens(response)
    ai_res = "".join(res)
    # 把模型结果也保存到chat.txt文件中
    sample_file.write(f"AI回复内容:{ai_res}\n")
    # 返回响应结果
    return ai_res

if __name__ == '__main__':
    print('您好,我是你的健康助手小黑:')
    while True:
        user_input = input('请您输入您的问题(q退出):')
        if user_input == 'q':
            print('欢迎下次再来')
            break
        # 模型预测获取结果
        my_model_res = model_predict(user_input)
        print(f"AI助手回复:{my_model_res}")


